import gymnasium as gym
import pygame
import numpy as np


def get_keyboard_action():
    """获取键盘输入并转换为游戏动作"""
    # 初始化动作 [转向, 油门, 刹车]
    action = np.array([0.0, 0.0, 0.0])

    # 获取当前按键状态
    keys = pygame.key.get_pressed()

    # 转向控制
    if keys[pygame.K_LEFT] or keys[pygame.K_a]:
        action[0] = -1.0  # 左转
    elif keys[pygame.K_RIGHT] or keys[pygame.K_d]:
        action[0] = 1.0  # 右转

    # 油门控制
    if keys[pygame.K_UP] or keys[pygame.K_w]:
        action[1] = 1.0  # 加速

    # 刹车控制
    if keys[pygame.K_DOWN] or keys[pygame.K_s]:
        action[2] = 1.0  # 刹车

    return action


def main():
    # 初始化pygame（用于键盘输入）
    pygame.init()

    # 创建环境
    env = gym.make("CarRacing-v2", render_mode="human")

    # 重置环境
    obs, info = env.reset()

    print("控制说明：")
    print("方向键 或 WASD 控制汽车")
    print("↑/W: 加速")
    print("↓/S: 刹车")
    print("←/A: 左转")
    print("→/D: 右转")
    print("ESC: 退出游戏")
    print("按任意键开始...")

    clock = pygame.time.Clock()
    running = True

    try:
        while running:
            # 处理pygame事件
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    running = False
                elif event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_ESCAPE:
                        running = False

            # 获取键盘动作
            action = get_keyboard_action()

            # 执行动作
            obs, reward, terminated, truncated, info = env.step(action)

            # 如果游戏结束，重新开始
            if terminated or truncated:
                print(f"游戏结束! 奖励: {reward:.2f}")
                obs, info = env.reset()

            # 控制帧率
            clock.tick(60)  # 60 FPS

    except KeyboardInterrupt:
        print("游戏被用户中断")

    finally:
        env.close()
        pygame.quit()


if __name__ == "__main__":
    main()
